import numpy as np
from tqdm import tqdm
from transformers import BertTokenizer

from config import Bert_Config

# 1. 载入 BERT 分词器
config = Bert_Config()
tokenizer = BertTokenizer.from_pretrained(config.bert_path)

# 2. 读取数据
data_path = config.train_path

lengths = []
invalid_lines = []

# 先统计总行数，方便进度条
with open(data_path, 'r', encoding='utf-8') as f:
    total_lines = sum(1 for _ in f)

# 正式读取并处理
with open(data_path, 'r', encoding='utf-8') as f:
    for line in tqdm(f, total=total_lines, desc="数据分析中..."):
        line = line.strip()

        if not line:
            invalid_lines.append(f"空行: {line}")
            continue

        try:
            text, label = line.split('\t')
        except ValueError:
            invalid_lines.append(f'格式有误忽略行: {line}')
            continue

        text, label = text.strip(), label.strip()

        # 检查文本和标签是否为空
        if not text or not label:
            invalid_lines.append(f"缺失文本或标签: {line}")
            continue

        # 检查标签是否合法数字
        if not label.isdigit():
            invalid_lines.append(f"非法标签: {line}")
            continue

        # 分词统计 token 长度（用 encode 方便统计真实长度，包含 [CLS] [SEP]）
        input_ids = tokenizer.encode(text)
        lengths.append(len(input_ids))

# 4. 转成 numpy 方便统计
lengths_np = np.array(lengths)

# 5. 统计结果
print(f"\n合法样本数: {len(lengths_np)}")
print(f"最大长度: {np.max(lengths_np)}")
print(f"最小长度: {np.min(lengths_np)}")
print(f"平均长度: {np.mean(lengths_np):.2f}")
print(f"中位数: {np.median(lengths_np)}")
print(f"标准差: {np.std(lengths_np):.2f}")
print(f"90分位数: {np.percentile(lengths_np, 90)}")
print(f"95分位数: {np.percentile(lengths_np, 95)}")
print(f"99分位数: {np.percentile(lengths_np, 99)}")

# 6. 打印异常行
print(f"\n异常行数: {len(invalid_lines)}")
for err in invalid_lines[:10]:  # 只展示前10行
    print(err)

# 打印结果
# 合法样本数: 180000
# 最大长度: 35      
# 最小长度: 4       
# 平均长度: 19.80   
# 中位数: 20.0  
# 标准差: 3.35  
# 90分位数: 24.0
# 95分位数: 25.0
# 99分位数: 26.0

# 异常行数: 0